#!/usr/bin/python
# -*- coding: utf-8 -*-
# @Version        : 1.0
# @Update Time    : 2025/4/8 21:28
# @File           : csv_utils.py
# @IDE            : PyCharm
# @Desc           : CSV导出工具类，支持大数据量导出和自定义编码

import csv
from io import StringIO, BytesIO, TextIOWrapper
from typing import List, Dict, Generator, Union, Optional, Iterable, Any, Iterator
from contextlib import contextmanager


class CSVExportCore:
    """CSV导出核心工具类，支持大数据量导出和自定义编码

    Attributes:
        headers: 表头列表
        summary_rows: 汇总行列表
        encoding: 文件编码
        include_bom: 是否包含BOM头
    """

    def __init__(
        self,
        headers: List[str],
        summary_rows: Optional[List[Dict[str, Any]]] = None,
        encoding: str = "utf-8-sig",
        include_bom: bool = True,
    ):
        """初始化CSV导出器

        Args:
            headers: 必须显式指定的表头
            summary_rows: 汇总行列表
            encoding: 文件编码，默认utf-8-sig
            include_bom: 是否包含BOM头，默认True
        """
        self.headers = headers
        self.summary_rows = summary_rows or []
        self.encoding = encoding
        self.include_bom = include_bom

    @contextmanager
    def _get_csv_writer(self, buffer: StringIO) -> Iterator[csv.DictWriter]:
        """获取CSV写入器的上下文管理器

        Args:
            buffer: 字符串缓冲区

        Yields:
            csv.DictWriter: CSV写入器
        """
        writer = csv.DictWriter(buffer, fieldnames=self.headers)
        try:
            yield writer
        finally:
            buffer.seek(0)
            buffer.truncate(0)

    def _generate_header(self) -> Generator[bytes, None, None]:
        """生成CSV头部

        Yields:
            bytes: 编码后的CSV头部数据
        """
        buffer = StringIO()
        with self._get_csv_writer(buffer) as writer:
            writer.writerow({col: col for col in self.headers})
            yield buffer.getvalue().encode(self.encoding)

    def _generate_summary(self) -> Generator[bytes, None, None]:
        """生成汇总行

        Yields:
            bytes: 编码后的汇总行数据
        """
        buffer = StringIO()
        with self._get_csv_writer(buffer) as writer:
            for row in self.summary_rows:
                writer.writerow(row)
                yield buffer.getvalue().encode(self.encoding)

    def full_data_stream(
        self,
        data: Union[List[Dict[str, Any]], Iterable[Dict[str, Any]]],
    ) -> Generator[bytes, None, None]:
        """生成CSV字节流

        Args:
            data: 支持列表或任意可迭代数据集

        Yields:
            bytes: 编码后的CSV数据流
        """
        # 统一转换为迭代器
        data_iter = iter(data)

        # 写入BOM头
        if self.include_bom:
            yield "\ufeff".encode(self.encoding)

        # 写入表头
        yield from self._generate_header()

        # 写入数据
        buffer = StringIO()
        with self._get_csv_writer(buffer) as writer:
            try:
                for row in data_iter:
                    writer.writerow(row)
                    yield buffer.getvalue().encode(self.encoding)
            finally:
                buffer.close()

        # 写入汇总行
        yield from self._generate_summary()

    def generator_stream(
        self, data_generator: Generator[Dict[str, Any], None, None]
    ) -> Generator[bytes, None, None]:
        """生成器数据流

        Args:
            data_generator: 数据生成器

        Yields:
            bytes: 编码后的CSV数据流
        """
        # 写入BOM头
        if self.include_bom:
            yield "\ufeff".encode(self.encoding)

        # 写入表头
        yield from self._generate_header()

        # 写入数据
        buffer = StringIO()
        with self._get_csv_writer(buffer) as writer:
            for row in data_generator:
                writer.writerow(row)
                yield buffer.getvalue().encode(self.encoding)

        # 写入汇总行
        yield from self._generate_summary()

    def to_bytes(
        self, data: Union[List[Dict[str, Any]], Iterable[Dict[str, Any]]]
    ) -> bytes:
        """生成完整字节数据

        Args:
            data: 支持列表或可迭代对象

        Returns:
            bytes: 完整的CSV文件字节数据
        """
        buffer = BytesIO()
        data_iter = iter(data)

        # 写入BOM头
        if self.include_bom:
            buffer.write("\ufeff".encode(self.encoding))

        # 使用文本包装器确保编码正确
        text_buffer = TextIOWrapper(
            buffer, encoding=self.encoding, newline="", write_through=True
        )

        writer = csv.DictWriter(text_buffer, fieldnames=self.headers)
        writer.writeheader()

        try:
            # 批量写入数据
            writer.writerows(data_iter)
            writer.writerows(self.summary_rows)
        finally:
            text_buffer.detach()  # 防止自动关闭底层buffer

        return buffer.getvalue()

    def to_file(
        self,
        data: Union[List[Dict[str, Any]], Iterable[Dict[str, Any]]],
        file_path: str,
    ) -> None:
        """保存到本地文件

        Args:
            data: 支持列表或可迭代对象
            file_path: 文件保存路径
        """
        data_iter = iter(data)

        with open(file_path, "wb") as f:
            # 写入BOM头
            if self.include_bom:
                f.write("\ufeff".encode(self.encoding))

            # 使用文本包装器确保编码正确
            text_wrapper = TextIOWrapper(
                f, encoding=self.encoding, newline="", write_through=True
            )

            writer = csv.DictWriter(text_wrapper, fieldnames=self.headers)
            writer.writeheader()

            try:
                # 批量写入数据
                writer.writerows(data_iter)
                writer.writerows(self.summary_rows)
            finally:
                text_wrapper.detach()  # 防止自动关闭底层文件对象
